import random
import math

def assign_cluster(x, c):
    min_dist = float('inf')  # 初始最小距离设为无穷大
    cluster_idx = 0          # 初始聚类索引设为0
    for i, centroid in enumerate(c):
        # 计算x与当前质心的欧氏距离（多维向量距离公式）
        dist = 0.0
        for xi, ci in zip(x, centroid):
            dist += (xi - ci) ** 2
        dist = math.sqrt(dist)
        # 更新最小距离和对应的聚类索引
        if dist < min_dist:
            min_dist = dist
            cluster_idx = i
    return cluster_idx
    
def Kmeans(data, k, epsilon=1e-3, iteration=100):
    if not isinstance(k, int) or k <= 0:
        raise ValueError("k必须是正整数")
    if len(data) <= k:
        raise ValueError("样本数量必须大于k")
    # 校验所有样本维度一致
    dim = len(data[0])
    for x in data:
        if len(x) != dim:
            raise ValueError("所有样本必须具有相同的维度")
    # 用索引随机选择，避免重复选择同一个样本
    random.seed(42)  # 固定随机种子，保证结果可复现（可删除）
    centroid_indices = random.sample(range(len(data)), k)
    centroids = [data[idx].copy() for idx in centroid_indices]
    for iter_cnt in range(iteration):
        # 保存当前质心（用于后续计算变化量）
        old_centroids = [c.copy() for c in centroids]
        cluster_assignments = []
        for x in data:
            cluster_idx = assign_cluster(x, centroids)
            cluster_assignments.append(cluster_idx)
        # 初始化：每个聚类的样本总和、样本数量
        cluster_sums = [[0.0 for _ in range(dim)] for _ in range(k)]  # 按维度求和
        cluster_counts = [0 for _ in range(k)]                        # 样本数量
        # 累加每个聚类的样本特征
        for idx, x in enumerate(data):
            cluster_idx = cluster_assignments[idx]
            for d in range(dim):
                cluster_sums[cluster_idx][d] += x[d]
            cluster_counts[cluster_idx] += 1
        # 计算新质心（均值）：避免除零（理论上不会发生，因k<=样本数且分配均匀）
        for i in range(k):
            if cluster_counts[i] == 0:
                # 极端情况：某个聚类无样本，重新随机选择一个样本作为质心
                centroids[i] = random.choice(data).copy()
            else:
                # 按维度计算均值
                for d in range(dim):
                    centroids[i][d] = cluster_sums[i][d] / cluster_counts[i]
        # 计算所有质心的最大变化量（多维向量的欧氏距离）
        max_centroid_change = 0.0
        for old_c, new_c in zip(old_centroids, centroids):
            change = 0.0
            for oc, nc in zip(old_c, new_c):
                change += (oc - nc) ** 2
            change = math.sqrt(change)
            if change > max_centroid_change:
                max_centroid_change = change
        # 若质心变化量小于epsilon，收敛并退出迭代
        if max_centroid_change < epsilon:
            print(f"迭代{iter_cnt+1}次后收敛（质心最大变化量：{max_centroid_change:.6f} < {epsilon}）")
            break
    return cluster_assignments, centroids

if __name__ == "__main__":
    data = [
        [1, 2], [2, 1], [1, 1], [2, 2],
        [5, 6], [6, 5], [5, 5], [6, 6],
        [9, 8], [8, 9], [9, 9], [8, 8]
    ]
    cluster_results, final_centroids = Kmeans(data, k=3)
    print("\n最终聚类结果（样本索引->聚类索引）：")
    for idx, cluster_idx in enumerate(cluster_results):
        print(f"样本{data[idx]} -> 聚类{cluster_idx}")
    print("\n最终质心：")
    for i, centroid in enumerate(final_centroids):
        print(f"聚类{i}质心：{[round(c, 3) for c in centroid]}")